import json
import sglang as sgl
import re
import time
from typing import Dict
from data_utils import load_chapter_main, load_doc_structure, match_chapter_ids_in_response, hard_coded_chapter_ids,read_sample_code
import os
import shutil
from logger import init_logger,add_file_handler,switch_handler
logger = init_logger(__name__)

# 环境变量设置
TASKS_FILE = "tasks.json"  # 问题文件
DOC_STRUCTURE_FILE = "doc_structure_remove_section_7_2.json"  # 文档结构文件
CHAPTERS_DIR = "./chapter_contents/"  # 章节文件目录（默认文件名为1.json, 2.json, 3.json, etc.）
EXAMPLE_ASCEND_C_FILE_DIR = "static_prompts"
HARD_CODED_CHAPTER_IDS = ["16.1.4.3.1","16.1.4.4.1"]
OUTPUT_FILE_PREFIX = "./generated_ascend_codes/codes"  # 输出文件
LOG_DIR = "./generated_ascend_codes/logs/"  # 输出文件
SAMPLE_CODE_TASK_DESCRIPTION = "sinh算子接收一个输入张量input，计算过程为对输入张量中的每个元素应用双曲正弦函数：sinh(x) = (e^x - e^(-x))/2。该算子执行逐元素操作，输出张量与输入张量具有完全相同的形状。双曲正弦函数在神经网络中可作为激活函数，在处理具有指数增长特性的数据时特别有用。请实现数值稳定且高效的sinh算子代码，确保处理大值时不会发生溢出。"

MAX_RETRIES = 3  # 最大重试次数
RETRY_DELAY = 5  # 重试延迟(秒)

REPEAT_TIMES = 1

API = "deepseek-r1"

if API == "deepseek-r1":
    API_MODEL_NAME = "deepseek-reasoner"
    API_KEY = "sk-c80e688b2d544f16825ceca256d3c6e2"  # API密钥
    API_BASE_URL = "https://api.deepseek.com"  # API基础URL
    STREAM = False
elif API == "huoshan-deepseek-r1":
    API_MODEL_NAME = "deepseek-r1-250120"
    API_KEY = "9b1afd3f-bbe8-493f-9ba4-c17784556f0c"  # API密钥
    API_BASE_URL = "https://ark.cn-beijing.volces.com/api/v3"  # API基础URL
    STREAM = False
elif API == "huoshan-deepseek-r1-distill-qwen-7b":
    API_MODEL_NAME = "deepseek-r1-distill-qwen-7b-250120"
    API_KEY = "9b1afd3f-bbe8-493f-9ba4-c17784556f0c"  # API密钥
    API_BASE_URL = "https://ark.cn-beijing.volces.com/api/v3"  # API基础URL
    STREAM = False
elif API == "qwq-official":
    API_MODEL_NAME = "qwq-32b"
    API_KEY = "sk-8db036d3267f461bb042c9c9b14d8143"
    API_BASE_URL = "https://dashscope.aliyuncs.com/compatible-mode/v1"
    STREAM = True


# 设置API后端
def setup_api():
    sgl.set_default_backend(
        sgl.OpenAI(
            API_MODEL_NAME,
            base_url=API_BASE_URL,
            api_key=API_KEY,
            is_chat_model=True
        )
    )
    logger.info(f"API = {API}")


# 加载任务问题
def load_tasks(tasks_file: str) -> Dict[str, str]:
    with open(tasks_file, 'r', encoding='utf-8') as f:
        return json.load(f)


# 查找与任务相关的章节，带异常重试
def find_related_sections_with_retry(sections_info, task_description, sample_code) -> str:
    for attempt in range(MAX_RETRIES):
        try:
            @sgl.function
            def find_related_sections(s):
                s += sgl.system(
                    f"你是一个精确的Ascend C编程文档检索助手。你的任务是从给定章节列表中找出与特定编程任务最相关的章节，严格避免引用不存在的章节。\n\n"
                    f"下面是Ascend C文档的完整章节列表及其摘要：\n{sections_info}\n\n"
                    f"请根据以下信息，找出完成给定任务最需要参考的章节：\n\n"
                    f"## 样例任务描述\n{SAMPLE_CODE_TASK_DESCRIPTION}\n\n"
                    f"请注意：\n"
                    f"1. 只能推荐章节列表中实际存在的章节\n"
                    f"2. 不要编造或推测任何未列出的章节\n"
                    f"3. 样例代码仅用于展示Ascend C的基本格式和语法特性，其解决的问题与当前任务不同\n"
                    f"4. 重点关注与当前任务直接相关的Ascend C专有功能、算法实现、优化技术和硬件接口的章节\n"
                    f"5. 如果任务与任何章节都不相关，请直接说明\"没有找到相关章节\"\n\n"
                    f"请按以下格式严格回答：\n"
                    f"\\chap{{实际章节号}}\\chap{{实际章节号}}\n"
                    f"简要解释这些章节如何帮助完成当前Ascend C编程任务（2-3句话）\n\n"
                )
                s += sgl.user(f"请找出与以下任务最相关的Ascend C文档章节：\"{task_description}\"。")
                s += sgl.assistant(sgl.gen(
                    name="result",
                    max_tokens=512,
                    temperature=0.7
                ))

            if STREAM:
                state = find_related_sections.run(stream=True)
            else:
                state = find_related_sections.run()
            logger.info(f"find_related_sections_with_retry message:{json.dumps(state.messages(), ensure_ascii=False, indent=4)}")
            return state
        except Exception as e:
            if attempt < MAX_RETRIES - 1:
                logger.info(f"章节匹配失败，错误信息: {e}，将在 {RETRY_DELAY} 秒后重试...")
                time.sleep(RETRY_DELAY)
            else:
                logger.info(f"章节匹配失败 {MAX_RETRIES} 次，使用默认章节。错误信息: {e}")
                return "\\chap{1} 由于匹配失败，默认使用Ascend C简介章节"


# 生成实现指定任务的 Ascend C 代码，带异常重试
def generate_ascend_c_codes_with_retry(task_description, sample_code, chapter_contents) -> str:
    for attempt in range(MAX_RETRIES):
        try:
            @sgl.function
            def generate_ascend_c_codes(s):
                s += sgl.system(
                    f"你是一个精通Ascend C编程的专家。你的任务是根据提供的文档内容生成符合特定任务需求的Ascend C代码。\n\n"
                    f"请仔细阅读以下内容：\n\n"
                    f"## 样例任务描述\n{SAMPLE_CODE_TASK_DESCRIPTION}\n\n"
                    f"## 样例Ascend C代码（仅供参考格式和语法）\n```cpp\n{sample_code}\n```\n\n"
                    f"## 相关文档内容\n{chapter_contents}\n\n"
                    f"请根据以上信息，生成一个完整且符合Ascend C规范的代码实现。确保代码：\n"
                    f"1. 遵循Ascend C的语法和最佳实践\n"
                    f"2. 满足任务描述中的所有要求\n"
                    f"3. 包含必要的注释解释代码的关键部分\n"
                    f"4. 完整可编译，不缺少关键组件\n"
                    f"5. 请你在写代码的时候一定要注意不要使用未定义的变量，同时注意变量的使用顺序\n"
                    f"6. 参考样例代码的格式和语法风格，但解决方案应针对当前任务定制\n\n"
                    f"只需提供代码，不需要额外解释。\n"
                )
                s += sgl.user(f"请根据上述文档内容和相关要求，针对以下任务生成Ascend C代码：{task_description}")
                s += sgl.assistant(sgl.gen(
                    name="result",
                    max_tokens=2048,
                    temperature=0.7
                ))

            if STREAM:
                state = generate_ascend_c_codes.run(stream=True)
            else:
                state = generate_ascend_c_codes.run()
            logger.info(f"generate_ascend_c_codes_with_retry message:{json.dumps(state.messages(),ensure_ascii=False, indent=4)}")
            return state
        except Exception as e:
            if attempt < MAX_RETRIES - 1:
                logger.info(f"章节匹配失败，错误信息: {e}，将在 {RETRY_DELAY} 秒后重试...")
                time.sleep(RETRY_DELAY)
            else:
                logger.info(f"章节匹配失败 {MAX_RETRIES} 次，使用默认章节。错误信息: {e}")
                return "\\chap{1} 由于匹配失败，默认使用Ascend C简介章节"


def ascend_code_generation():
    """
    主函数：根据文档目录、样例代码和任务描述生成Ascend C代码
    返回：生成的Ascend C代码
    """
    # 加载任务
    logger.info("Loading tasks...")
    tasks = load_tasks(TASKS_FILE)

    # 加载文档结构
    logger.info("Loading doc structures...")
    sections_info = load_doc_structure(DOC_STRUCTURE_FILE)

    # 读取样例代码
    sample_code = read_sample_code(EXAMPLE_ASCEND_C_FILE_DIR)
    if sample_code is None:
        return "无法读取样例代码文件。"

    # 处理每个任务
    logger.info("start processing codegen tasks ...")
    total_tasks = len(tasks)
    for idx,(task_id, task) in enumerate(tasks.items()):
        #每一个task都会输出到单独的文件
        if idx ==0:
            old_handler = add_file_handler(os.path.join(LOG_DIR,f"{API}_task_{task_id}.log"))
        else:
            old_handler = switch_handler(os.path.join(LOG_DIR,f"{API}_task_{task_id}.log"), old_handler)
        # handler = RotatingFileHandler(filename=)
        # logging.getLogger().addHandler(handler)
        logger.info(f"处理任务 {idx}/{total_tasks}: {task[:50]}...")

        # 步骤1: 找到相关章节
        logger.info(f"try to find related sections...")
        related_sections_response = find_related_sections_with_retry(sections_info, task, sample_code)
        related_chapter_ids = match_chapter_ids_in_response(related_sections_response["result"], task_id)
        related_chapter_ids = hard_coded_chapter_ids(related_chapter_ids,HARD_CODED_CHAPTER_IDS)
        # related_chapter_ids = [ '6.2.1.2', '5.1', '6.2.3', '7.1', '16.1.4.2.8.3', '16.1.4.3.1', '16.1.4.4.1']
        logger.info(f"related_chapter_ids = {related_chapter_ids}")

        # 步骤2: 加载相关章节内容
        logger.info(f"try to load related sections...")
        all_chapter_contents = load_chapter_main(related_chapter_ids, chapters_dir=CHAPTERS_DIR)

        # 步骤3: 生成Ascend C代码
        logger.info(f"try to generate ascend c codes...")

        for i in range(REPEAT_TIMES):
            logger.info(f"{i}th attempt begins...")
            generated_codes = generate_ascend_c_codes_with_retry(task, sample_code, all_chapter_contents)["result"]
            # generated_codes = re.sub(r'^```cpp\n|\n```$', '', generated_codes)
            generated_codes = re.sub(r'```cpp\s*\n(.*?)\n\s*```.*', r'\1', generated_codes, flags=re.DOTALL)
            if not generated_codes:
                logger.info("No cpp file generated.")
                continue
            if generated_codes:
                with open(os.path.join(f"{OUTPUT_FILE_PREFIX}",f"{API}_task_{task_id}_{i}.cpp"), "w", encoding="utf-8") as f:
                    f.write(generated_codes)
            logger.info(f"{i}th attempt successes...")

    return generated_codes


# 示例用法
if __name__ == "__main__":
    os.makedirs(OUTPUT_FILE_PREFIX, exist_ok=True)
    os.makedirs(LOG_DIR, exist_ok=True)
    # logging.basicConfig(filename=os.path.join(LOG_DIR,'codegen_default.log'), level=logging.INFO)
    setup_api()

    # 生成代码
    ascend_code_generation()
